## dehaze
import numpy as np
import cv2,math

def darkChannel(p, winSize):
    '''
    darkChannel就是先取rgb通道的最小值然后再进行一次最小值滤波
    '''
    minChannel = p.min(axis=2)
    
    shape = cv2.MORPH_RECT
    kernel = cv2.getStructuringElement(shape, winSize)
    min_image = cv2.erode(minChannel, kernel)
    return min_image

def guideFilter(I, p, winSize, eps):
    """
    导向图像(Guidance Image) I，滤波输出图像(Filtering Input Image) p，均值平滑窗口半径 r，正则化参数 e。
    利用导向滤波进行图像平滑处理时，通常令p=I。
    其中：guideFilter(）函数调用opencv自带的库函数blur() 进行均值平滑。
    :param I:
    :param p:
    :param winSize:
    :param eps:
    :return:
    """
    
    sqaure_avg = cv2.blur(p * p, winSize)
    avg = cv2.blur(p, winSize)
    
    ak = (sqaure_avg - avg ** 2)/(sqaure_avg - avg **2 + eps)
    bk = (-1 * ak + 1) * avg
    
    a = cv2.blur(ak, winSize)
    b = cv2.blur(bk, winSize)
    
    q = a * p + b
    
    return q
def largest_indices(ary, n):
    """Returns the n largest indices from a numpy array."""
    flat = ary.flatten()
    indices = np.argpartition(flat, -n)[-n:]
    indices = indices[np.argsort(-flat[indices])]
    return np.unravel_index(indices, ary.shape)

def CalculateA(p, winSize):
    '''
    A的计算和暗通道图有关，找出前0.1%的像素，将他们最亮的通道点标出来然后，找到原图中对应像素点的值作为大气光A
    '''
    dark_p = darkChannel(p, winSize)
    
    h, w = dark_p.shape
    top_pixels = h * w // 1000
    
    indexs = largest_indices(dark_p, top_pixels)
    
    A = p[indexs].max(axis=0)
    
    return A

def CalculateTx(p, winSize,w=0.95):

    #这个函数有问题，最后没有用到，算是一次尝试
    A = CalculateA(p,winSize)
    
    dark_t = darkChannel(p / A,winSize)
    
    t = 1 - w * dark_t
    
    return t
    

def MistRemove(p, winSize,t0):
    
    A = CalculateA(p,winSize)
    
    p=p/255
    dark_p = darkChannel(p,winSize)
    
    #用暗通道图作引导图进行导向滤波
    b, g, r = cv2.split(p)
    gf1 = guideFilter(b, dark_p, winSize, math.pow(0.1,2))
    gf2 = guideFilter(g, dark_p, winSize, math.pow(0.1,2))
    gf3 = guideFilter(r, dark_p, winSize, math.pow(0.1,2))
    
    gf = cv2.merge([gf1, gf2, gf3])
    
    gf = gf*255
    gf[gf>255] = 255
    gf = np.round(gf)
    gf = gf.astype(np.uint8)
    
    t = 1 - gf * 0.95 / A #折射率的球阀，本来0.95应该是一个系数，但我们这里直接用0.95替代了，因为这是一个比较好的值
    
    
    
    t = cv2.max(t, t0) #最后要除以t，所以如果t太小会比较不好处理，所以这里给一个阈值
    
   
    

    rows, cols, channels = gf.shape
    
    a = np.average(A) #方便处理这里直接
    
    p = p*255
    p[p>255] = 255
    p = np.round(p)
    p = p.astype(np.uint8)
    
    o = (p -a)/t+a
    o = np.round(o)
    o[o>255]=255 #老师上课说可能没有处理超纲的值，确实，后来发现确实是这个问题，这行是后来加的
    o = o.astype("uint8")
    
    res = np.hstack((p,o))
    
    cv2.imshow('1',res)
    cv2.waitKey(0)
    
    return o
    
    


p = cv2.imread(input("Please input the file name\n"))
MistRemove(p,(16,16), 0.1)
